"""
示例 6: 创建复杂电极几何

演示如何：
1. 创建复杂的 3D 电极结构
2. 使用数学函数定义几何
3. 创建实用的离子光学元件
"""

import sys
import os
from math import sqrt, sin, cos, pi

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, project_root)

from SIMION.PA import PA


def create_tof_analyzer():
    """创建飞行时间质谱(TOF)分析器"""
    
    print("\n" + "="*60)
    print("创建飞行时间质谱分析器")
    print("="*60 + "\n")
    
    pa = PA(
        nx=80,
        ny=200,
        nz=1,
        symmetry='cylindrical',
        max_voltage=5000,
        field_type='electrostatic',
        dx_mm=0.5,
        dy_mm=0.5
    )
    
    print(f"创建阵列: {pa.nx()} x {pa.ny()} (圆柱对称)\n")
    
    # TOF 结构:
    # 1. 离子源区域（提取电极）
    # 2. 加速区域
    # 3. 漂移区域（无场）
    # 4. 反射器（可选）
    
    inner_radius = 5
    outer_radius = 70
    
    # 1. 提取电极 (Y=10-15, 0V)
    print("创建提取电极...")
    for y in range(10, 16):
        for r in range(inner_radius, outer_radius + 1):
            pa.point(r, y, 0, 1, 0.0)
    
    # 2. 加速电极 (Y=30-35, -2000V)
    print("创建加速电极...")
    for y in range(30, 36):
        for r in range(inner_radius, outer_radius + 1):
            pa.point(r, y, 0, 1, -2000.0)
    
    # 3. 漂移区域入口 (Y=50-52, -2000V)
    print("创建漂移区域入口...")
    for y in range(50, 53):
        for r in range(inner_radius, outer_radius + 1):
            pa.point(r, y, 0, 1, -2000.0)
    
    # 4. 反射器电极组（用于反射式TOF）
    # 多个电极形成减速场
    print("创建反射器电极...")
    reflector_start = 150
    num_reflector_electrodes = 5
    
    for i in range(num_reflector_electrodes):
        y_start = reflector_start + i * 8
        y_end = y_start + 3
        voltage = -2000.0 + i * 500.0  # 电压逐渐升高
        
        for y in range(y_start, y_end + 1):
            for r in range(inner_radius, outer_radius + 1):
                pa.point(r, y, 0, 1, voltage)
    
    # 保存文件
    output_file = os.path.join(project_root, "SIMION", "examples", "pa", "tof_analyzer.pa#")
    print(f"\n保存到: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("创建完成!")
    print("="*60)
    print("\nTOF 分析器结构:")
    print("  - 提取区域: Y=10-15, 0V")
    print("  - 加速区域: Y=30-35, -2000V")
    print("  - 漂移区域: Y=50-150 (无场)")
    print("  - 反射器: Y=150-190, -2000V 到 0V\n")
    
    return pa


def create_orbitrap_like():
    """创建类似 Orbitrap 的电极几何"""
    
    print("\n" + "="*60)
    print("创建类似 Orbitrap 的电极几何")
    print("="*60 + "\n")
    
    pa = PA(
        nx=100,
        ny=200,
        nz=1,
        symmetry='cylindrical',
        max_voltage=10000,
        field_type='electrostatic',
        dx_mm=0.25,
        dy_mm=0.25
    )
    
    print(f"创建阵列: {pa.nx()} x {pa.ny()} (圆柱对称)\n")
    
    # Orbitrap 使用抛物线形电极
    # 中心电极: 细长的梭形
    # 外电极: 抛物线形
    
    center_y = pa.ny() // 2
    
    print("创建中心电极（梭形）...")
    
    # 中心电极半径随轴向位置变化
    for y in range(pa.ny()):
        # 距离中心的距离
        dy = abs(y - center_y)
        
        # 梭形半径: 中间细，两端粗
        if dy < 80:
            # 使用椭球方程
            center_radius = 5 + int(10 * sqrt(1 - (dy/80)**2))
            
            for r in range(0, center_radius + 1):
                pa.point(r, y, 0, 1, 1000.0)
    
    print("创建外电极（抛物线形）...")
    
    # 外电极: 抛物线形
    k = 0.02  # 抛物线系数
    
    for y in range(pa.ny()):
        dy = y - center_y
        
        # 抛物线方程: r = r0 + k*y^2
        outer_radius = 50 + int(k * dy**2)
        
        if outer_radius < pa.nx():
            # 电极厚度
            for r in range(outer_radius, min(outer_radius + 5, pa.nx())):
                pa.point(r, y, 0, 1, 0.0)
    
    # 保存文件
    output_file = os.path.join(project_root, "SIMION", "examples", "pa", "orbitrap_like.pa#")
    print(f"\n保存到: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("创建完成!")
    print("="*60)
    print("\n这是一个简化的 Orbitrap 类型几何")
    print("中心电极: 1000V")
    print("外电极: 0V\n")
    
    return pa


def create_ion_trap():
    """创建 3D 保罗离子阱"""
    
    print("\n" + "="*60)
    print("创建 3D 保罗离子阱")
    print("="*60 + "\n")
    
    pa = PA(
        nx=80,
        ny=80,
        nz=80,
        symmetry='planar',
        max_voltage=1000,
        field_type='electrostatic',
        dx_mm=0.1,
        dy_mm=0.1,
        dz_mm=0.1
    )
    
    print(f"创建阵列: {pa.nx()} x {pa.ny()} x {pa.nz()} (3D)\n")
    
    center_x = pa.nx() // 2
    center_y = pa.ny() // 2
    center_z = pa.nz() // 2
    
    # 离子阱参数
    r0 = 10  # 阱的特征尺寸
    z0 = 14  # 端帽间距的一半
    
    print(f"离子阱参数:")
    print(f"  r0 = {r0} 网格")
    print(f"  z0 = {z0} 网格")
    print(f"  中心: ({center_x}, {center_y}, {center_z})\n")
    
    # 环形电极 (双曲线形)
    print("创建环形电极...")
    ring_count = 0
    
    for z in range(pa.nz()):
        for y in range(pa.ny()):
            for x in range(pa.nx()):
                dx = x - center_x
                dy = y - center_y
                dz = z - center_z
                
                r_sq = dx**2 + dy**2
                
                # 双曲面方程: r^2 - 2*z^2 = r0^2
                hyperbola_value = r_sq - 2 * dz**2
                
                # 环形电极区域
                if abs(hyperbola_value - r0**2) < 15 and abs(dz) < z0 - 5:
                    pa.point(x, y, z, 1, 100.0)
                    ring_count += 1
    
    print(f"环形电极点数: {ring_count}")
    
    # 端帽电极 (双曲线形)
    print("创建端帽电极...")
    endcap_count = 0
    
    for z in range(pa.nz()):
        for y in range(pa.ny()):
            for x in range(pa.nx()):
                dx = x - center_x
                dy = y - center_y
                dz = z - center_z
                
                r_sq = dx**2 + dy**2
                
                # 端帽双曲面: 2*z^2 - r^2 = z0^2
                endcap_value = 2 * dz**2 - r_sq
                
                # 两个端帽
                if abs(endcap_value - z0**2) < 20 and r_sq < (2*r0)**2:
                    pa.point(x, y, z, 1, 0.0)
                    endcap_count += 1
    
    print(f"端帽电极点数: {endcap_count}")
    
    # 验证中心
    is_elec, pot = pa.point(center_x, center_y, center_z)
    print(f"\n验证阱中心:")
    print(f"  坐标: ({center_x}, {center_y}, {center_z})")
    print(f"  是电极: {is_elec} (应该是 False)")
    
    # 保存文件
    output_file = os.path.join(project_root, "SIMION", "examples", "pa", "ion_trap_3d.pa#")
    print(f"\n保存到: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("创建完成!")
    print("="*60)
    print("\n3D 保罗离子阱:")
    print("  - 环形电极: 100V")
    print("  - 端帽电极: 0V")
    print("\n注意: 这是一个简化模型，实际的离子阱需要更精确的几何\n")
    
    return pa


def create_quadrupole_mass_filter():
    """创建四极杆质量过滤器"""
    
    print("\n" + "="*60)
    print("创建四极杆质量过滤器")
    print("="*60 + "\n")
    
    pa = PA(
        nx=60,
        ny=60,
        nz=120,
        symmetry='planar',
        max_voltage=500,
        field_type='electrostatic',
        dx_mm=0.2,
        dy_mm=0.2,
        dz_mm=0.5
    )
    
    print(f"创建阵列: {pa.nx()} x {pa.ny()} x {pa.nz()} (3D)\n")
    
    center_x = pa.nx() // 2
    center_y = pa.ny() // 2
    
    # 四极杆参数
    r0 = 15  # 场半径（从中心到杆的距离）
    rod_radius = 8  # 杆的半径
    
    print(f"四极杆参数:")
    print(f"  场半径 r0: {r0} 网格 ({r0*pa.dx_mm()} mm)")
    print(f"  杆半径: {rod_radius} 网格 ({rod_radius*pa.dx_mm()} mm)\n")
    
    # 四根杆的位置（45度旋转）
    rod_positions = [
        (center_x + r0, center_y, +100.0),  # 右
        (center_x - r0, center_y, -100.0),  # 左
        (center_x, center_y + r0, +100.0),  # 上
        (center_x, center_y - r0, -100.0),  # 下
    ]
    
    print("创建四根杆...")
    
    for i, (rod_x, rod_y, voltage) in enumerate(rod_positions):
        print(f"  杆 {i+1}: 位置=({rod_x},{rod_y}), 电压={voltage}V")
        
        rod_count = 0
        # 沿 Z 方向创建整根杆
        for z in range(pa.nz()):
            for y in range(pa.ny()):
                for x in range(pa.nx()):
                    # 到杆中心的距离
                    dist = sqrt((x - rod_x)**2 + (y - rod_y)**2)
                    
                    # 如果在杆的半径内
                    if dist <= rod_radius:
                        pa.point(x, y, z, 1, voltage)
                        rod_count += 1
        
        print(f"    电极点数: {rod_count}")
    
    # 验证中心轴
    is_elec, pot = pa.point(center_x, center_y, pa.nz()//2)
    print(f"\n验证中心轴:")
    print(f"  坐标: ({center_x}, {center_y}, {pa.nz()//2})")
    print(f"  是电极: {is_elec} (应该是 False)")
    
    # 保存文件
    output_file = os.path.join(project_root, "SIMION", "examples", "pa", "quadrupole_filter.pa#")
    print(f"\n保存到: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("创建完成!")
    print("="*60)
    print("\n四极杆质量过滤器:")
    print("  - 对角杆电压: +100V 和 -100V")
    print("  - 实际使用时需要叠加射频电压\n")
    
    return pa


if __name__ == "__main__":
    print("\n" + "="*60)
    print("SIMION 复杂电极几何创建示例")
    print("="*60)
    
    # 创建各种复杂几何
    print("\n这些示例展示了如何创建实用的离子光学元件")
    print("每个示例可能需要几秒到几十秒的时间来创建\n")
    
    # 1. TOF 分析器
    print("\n[1/4] 创建 TOF 分析器...")
    pa1 = create_tof_analyzer()
    
    # 2. Orbitrap 类型几何
    print("\n[2/4] 创建 Orbitrap 类型几何...")
    pa2 = create_orbitrap_like()
    
    # 3. 3D 离子阱
    print("\n[3/4] 创建 3D 离子阱...")
    pa3 = create_ion_trap()
    
    # 4. 四极杆质量过滤器
    print("\n[4/4] 创建四极杆质量过滤器...")
    pa4 = create_quadrupole_mass_filter()
    
    print("\n" + "="*60)
    print("所有复杂几何创建完成!")
    print("="*60)
    print("\n生成的文件:")
    print("  - tof_analyzer.pa# : 飞行时间质谱分析器")
    print("  - orbitrap_like.pa# : Orbitrap 类型几何")
    print("  - ion_trap_3d.pa# : 3D 保罗离子阱")
    print("  - quadrupole_filter.pa# : 四极杆质量过滤器")
    print("\n可以在 SIMION 中打开这些文件并进行 Refine")
    print("然后可以进行粒子飞行模拟\n")

